// Copyright (c) 2023 Presto Labs Pte. Ltd.
// Author: leon, junxiao

#include <cmath>

#include "coin2/strategy/dex_strat/dex_strat_utils.h"

namespace coin2::strategy::dex_strat {



WindowEstimator::WindowEstimator(int64_t max_record_cnt, double range_multiplier)
  : max_record_cnt_(max_record_cnt), range_multiplier_(range_multiplier) {
    cur_record_cnt_ = 0;
    cur_mean_ = 0.0;
    cur_sqr_mean_ = 0.0;
    half_win_size_ = 0.0;
}

bool WindowEstimator::Empty() {
  return cur_record_cnt_ <= 0;
}

bool WindowEstimator::HasEnoughRecord() {
  return cur_record_cnt_ >= max_record_cnt_;
}

double WindowEstimator::GetAvgPrice() {
  return cur_mean_;
}

double WindowEstimator::GetHalfWinSize() {
  return half_win_size_;
}

std::pair<double, double> WindowEstimator::GetPriceWindow() {
  return {cur_mean_ - half_win_size_, cur_mean_ + half_win_size_};
}

void WindowEstimator::InsertRecord(double price) {
  if (cur_record_cnt_ < max_record_cnt_) {
    record_queue_.push_back(price);
    cur_mean_ = cur_mean_ * (static_cast<double>(cur_record_cnt_) / static_cast<double>(cur_record_cnt_ + 1))
      + price / static_cast<double>(cur_record_cnt_ + 1);
    cur_sqr_mean_ = cur_sqr_mean_ * (static_cast<double>(cur_record_cnt_) / static_cast<double>(cur_record_cnt_ + 1))
      + (price * price) / static_cast<double>(cur_record_cnt_ + 1);
    cur_record_cnt_++;
  } else {
    record_queue_.push_back(price);
    double poped_price = record_queue_.front();
    record_queue_.pop_front();
    cur_mean_ += (price - poped_price) / static_cast<double>(cur_record_cnt_);
    cur_sqr_mean_ += (price * price - poped_price * poped_price) / static_cast<double>(cur_record_cnt_);
  }
  half_win_size_ = std::sqrt(std::max(0.0, cur_sqr_mean_ - cur_mean_ * cur_mean_)) * range_multiplier_;
}

void WindowEstimator::PrintInfo() {
  LOG(INFO) << fmt::format(
    "record_cnt:{}, cur_mean:{}, cur_sqr_mean:{}, price_low:{}, price_upper:{}",
    cur_record_cnt_, cur_mean_, cur_sqr_mean_, cur_mean_ - half_win_size_, cur_mean_ + half_win_size_
  );
}

BiasCalculator::BiasCalculator(int64_t max_record_cnt, double withdraw_signal_bar, double provide_signal_bar)
  : max_record_cnt_(max_record_cnt), withdraw_signal_bar_(withdraw_signal_bar), provide_signal_bar_(provide_signal_bar) {
  cur_record_cnt_ = 0;
  cur_bias_val_ = 0.0;
}

double BiasCalculator::CalculateBias(double price_diff) {
  if (price_diff > 0.0) {
    return price_diff * price_diff;
  } else {
    return -(price_diff * price_diff);
  }
}

bool BiasCalculator::Empty() {
  return cur_record_cnt_ <= 0;
}

bool BiasCalculator::HasEnoughRecord() {
  return cur_record_cnt_ >= max_record_cnt_;
}

double BiasCalculator::GetAvgBiasVal() {
  return cur_bias_val_;
}

double BiasCalculator::GetLastBiasVal() {
  if (Empty()) {
    return 0.0;
  }
  return CalculateBias(record_queue_.back());
}

bool BiasCalculator::IsPlusBiasWithdraw() {
  return cur_bias_val_ > withdraw_signal_bar_;
}

bool BiasCalculator::IsMinusBiasWithdraw() {
  return cur_bias_val_ < -withdraw_signal_bar_;
}

bool BiasCalculator::IsPlusBiasProvide() {
  return cur_bias_val_ > provide_signal_bar_;
}

bool BiasCalculator::IsMinusBiasProvide() {
  return cur_bias_val_ < -provide_signal_bar_;
}

bool BiasCalculator::ShouldWithdraw() {
  //LOG(INFO) << "BiasCalculator::ShouldWithdraw";
  return IsPlusBiasWithdraw() && IsMinusBiasWithdraw();
}

bool BiasCalculator::ShouldProvide() {
  //LOG(INFO) << "BiasCalculator::ShouldProvide";
  return !IsPlusBiasProvide() && !IsMinusBiasProvide();
}

void BiasCalculator::InsertRecord(double trade_price, double ref_price) {
  double price_diff = ref_price - trade_price;
  if (cur_record_cnt_ < max_record_cnt_) {
    record_queue_.push_back(price_diff);
    cur_bias_val_ = cur_bias_val_ * (static_cast<double>(cur_record_cnt_) / static_cast<double>(cur_record_cnt_ + 1))
      + CalculateBias(price_diff) / static_cast<double>(cur_record_cnt_ + 1);
    cur_record_cnt_++;
  } else {
    record_queue_.push_back(price_diff);
    double poped_price_diff = record_queue_.front();
    record_queue_.pop_front();
    cur_bias_val_ += (CalculateBias(price_diff) - CalculateBias(poped_price_diff)) / static_cast<double>(cur_record_cnt_);
  }
}

void BiasCalculator::PrintInfo() {
  double bias_level = std::sqrt(std::abs(cur_bias_val_));
  LOG(INFO) << fmt::format(
    "record_cnt:{}, cur_bias_val:{}, bias_level:{}",
    cur_record_cnt_, cur_bias_val_, bias_level
  );
}

}  // namespace coin2::strategy::dex_strat
